from ultralytics import FastSAM
from ultralytics.models.fastsam import FastSAMPredictor

# Define an inference source
source = "/home/champrin/PycharmProjects/keypoints_train/Screenshot from 2025-02-04 23-41-53.png"

model_path = "/home/champrin/PycharmProjects/keypoints_train/FastSAM-s.pt"

# Create a FastSAM model
model = FastSAM(model_path)  # or FastSAM-x.pt

# Run inference on an image
everything_results = model(source, retina_masks=True, imgsz=640, conf=0.5, iou=0.01, save=False)

# Run inference with bboxes prompt
# results = model(source, bboxes=[439, 437, 524, 709], save=True)

# Run inference with points prompt
# results = model(source, points=[[200, 200]], labels=[1], save=False)

# Run inference with texts prompt
# results = model(source, texts="a photo of a dog", save=False)

# Run inference with bboxes and points and texts prompt at the same time
# results = model(source, bboxes=[439, 437, 524, 709], points=[[200, 200]], labels=[1], texts="a photo of a dog", save=False)


# Create FastSAMPredictor
# overrides = dict(conf=0.5, task="segment", mode="predict", model="/home/champrin/PycharmProjects/keypoints_train/FastSAM-s.pt", save=True, imgsz=1024)
# predictor = FastSAMPredictor(overrides=overrides)

# Segment everything
# everything_results = predictor(source)

# Prompt inference
# bbox_results = predictor.prompt(everything_results, bboxes=[[200, 200, 300, 300]])
# point_results = predictor.prompt(everything_results, points=[200, 200])
# text_results = predictor.prompt(everything_results, texts="a photo of a dog")